package shawnho.ideaagent;

import javassist.ClassPool;
import javassist.CtClass;
import javassist.bytecode.*;

import java.io.ByteArrayInputStream;
import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.IllegalClassFormatException;
import java.security.ProtectionDomain;
import java.util.List;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * @author shawnho
 */
public class IdeaTransformer implements ClassFileTransformer {
    private ClassPool pool;
    private AtomicInteger count;
    private boolean iedis;
    private boolean mybatis;

    public IdeaTransformer() {
        this.pool = ClassPool.getDefault();
        this.count = new AtomicInteger();
        this.iedis = false;
        this.mybatis = false;
    }

    private int enableAgent(ConstPool constPool, MethodInfo methodInfo, CodeIterator iter) throws BadBytecode {
        int i = 0;
        int[] code = new int[]{187, 89, 25, 183, 58, 1, 58, 3, 54, 21, 25, 185};
        int[] save = new int[code.length];
        iter.begin();
        while (iter.hasNext()) {
            int index = iter.next();
            int opcode = iter.byteAt(index);
            if (opcode == code[i]) {
                save[i] = index;
                if (++i == code.length) {
                    Bytecode bytecode = new Bytecode(constPool);
                    bytecode.addLdc(constPool.addStringInfo("-ideaagent:"));
                    bytecode.addAstore(iter.u16bitAt(save[4]) & 255);
                    iter.insert(save[5], bytecode.get());
                    return bytecode.getSize();
                }
            } else {
                i = 0;
            }
        }
        return 0;
    }

    private int modifyUrl(ConstPool constPool, MethodInfo methodInfo, CodeIterator iter) throws BadBytecode {
        int i = 0;
        int[] code = new int[]{187, 89, 183, 25, 182, 187, 89, 25, 183, 182, 182, 184, 4, 182, 4, 182};
        int[] save = new int[code.length];
        iter.begin();
        while (iter.hasNext()) {
            int index = iter.next();
            int opcode = iter.byteAt(index);
            if (opcode == code[i]) {
                save[i] = index;
                if (++i == code.length) {
                    Bytecode bytecode = new Bytecode(constPool);
                    bytecode.addOpcode(Opcode.POP);
                    bytecode.addLdc(constPool.addStringInfo("https://localhost:7848/q1"));
                    iter.insert(save[11], bytecode.get());
                    return bytecode.getSize();
                }
            } else {
                i = 0;
            }
        }
        return 0;
    }

    private int modifyCert(ConstPool constPool, MethodInfo methodInfo, CodeIterator iter) throws BadBytecode {
        int i = 0;
        int[] code = new int[]{187, 89, 187, 89, 183, 25, 182, 187, 89, 184, 183, 182, 25, 182, 182, 183, 182, 58};
        int[] save = new int[code.length];
        iter.begin();
        while (iter.hasNext()) {
            int index = iter.next();
            int opcode = iter.byteAt(index);
            if (opcode == code[i]) {
                save[i] = index;
                if (++i == code.length) {
                    Bytecode bytecode = new Bytecode(constPool);
                    bytecode.addOpcode(Opcode.POP);
                    bytecode.addLdc(constPool.addStringInfo("-----BEGIN CERTIFICATE-----\n" +
                            "MIICYDCCAcmgAwIBAgIEAdWdvDANBgkqhkiG9w0BAQsFADBjMQswCQYDVQQGEwJj\n" +
                            "bjELMAkGA1UECBMCc2MxCzAJBgNVBAcTAmNkMRIwEAYDVQQKEwlsb2NhbGhvc3Qx\n" +
                            "EjAQBgNVBAsTCWxvY2FsaG9zdDESMBAGA1UEAxMJbG9jYWxob3N0MB4XDTE4MDQw\n" +
                            "NjA5MDk0NVoXDTE5MDQwNjA5MDk0NVowYzELMAkGA1UEBhMCY24xCzAJBgNVBAgT\n" +
                            "AnNjMQswCQYDVQQHEwJjZDESMBAGA1UEChMJbG9jYWxob3N0MRIwEAYDVQQLEwls\n" +
                            "b2NhbGhvc3QxEjAQBgNVBAMTCWxvY2FsaG9zdDCBnzANBgkqhkiG9w0BAQEFAAOB\n" +
                            "jQAwgYkCgYEAi5LrPK3cFLME8Agqo4w1ZA/RcYUZ/CE8x8vU/MfWz7DmLn7GEayo\n" +
                            "kbFeYVoyXqdKi2rCPB+ipJw/KNo3WTsrn7eqfSnGxeh1xaPJvQr9lTMRGWhnxoQU\n" +
                            "CtVvNhjzicIdqQoTFZYerxb8RZofzgY3whr6w212Q1VufXcfTbzbJiUCAwEAAaMh\n" +
                            "MB8wHQYDVR0OBBYEFA6H5FhG+9GW4IlIdZvRmklsyXvVMA0GCSqGSIb3DQEBCwUA\n" +
                            "A4GBAH4rvj6Zh++zTG2ikF/zmO++cBUkcyV1bxys2QujaALVCxmLvVCpUaJ1TchS\n" +
                            "IYIJkF3weDCr71z5e7aTmCmsb5aL+li1Ycn4HLvsS5xjI12OkmgDmgjp71vSy70O\n" +
                            "CkLaadSjCnYI6x6lMc7UhuKROXEXKbs1DTxjyAkTnnPkXQeS\n" +
                            "-----END CERTIFICATE-----"));
                    iter.insert(save[15], bytecode.get());
                    return bytecode.getSize();
                }
            } else {
                i = 0;
            }
        }
        return 0;
    }

    private int modifyResult(ConstPool constPool, MethodInfo methodInfo, CodeIterator iter) throws BadBytecode {
        int i = 0;
        int[] code = new int[]{187, 89, 25, 178, 183, 184, 58, 184, 25, 3, 51};
        int[] save = new int[code.length];
        iter.begin();
        while (iter.hasNext()) {
            int index = iter.next();
            int opcode = iter.byteAt(index);
            if (opcode == code[i]) {
                save[i] = index;
                if (++i == code.length) {
                    Bytecode bytecode = new Bytecode(constPool);
                    bytecode.addInvokestatic("shawnho.ideaagent.Utils", "getObstructResult", "()[B");
                    bytecode.addAstore(iter.u16bitAt(save[6]) & 255);
                    iter.insert(save[7], bytecode.get());
                    return bytecode.getSize();
                }
            } else {
                i = 0;
            }
        }
        return 0;
    }

    private void modifyIedis(ConstPool constPool, MethodInfo methodInfo, CodeIterator iter) throws BadBytecode {
        if (methodInfo.getName().equals("b") && methodInfo.getDescriptor().equals("()V")) {
            Bytecode bytecode = new Bytecode(constPool);
            bytecode.addOpcode(Opcode.POP);
            bytecode.addLdc(constPool.addStringInfo("https://localhost:7848/q2"));
            iter.insert(17555, bytecode.get());
        } else if (methodInfo.getName().equals("b") && methodInfo.getDescriptor().equals("(Lcom/intellij/util/io/HttpRequests$Request;)V")) {
            Bytecode cert = new Bytecode(constPool);
            cert.addOpcode(Opcode.POP);
            cert.addLdc(constPool.addStringInfo("-----BEGIN CERTIFICATE-----\n" +
                    "MIICYDCCAcmgAwIBAgIEAdWdvDANBgkqhkiG9w0BAQsFADBjMQswCQYDVQQGEwJj\n" +
                    "bjELMAkGA1UECBMCc2MxCzAJBgNVBAcTAmNkMRIwEAYDVQQKEwlsb2NhbGhvc3Qx\n" +
                    "EjAQBgNVBAsTCWxvY2FsaG9zdDESMBAGA1UEAxMJbG9jYWxob3N0MB4XDTE4MDQw\n" +
                    "NjA5MDk0NVoXDTE5MDQwNjA5MDk0NVowYzELMAkGA1UEBhMCY24xCzAJBgNVBAgT\n" +
                    "AnNjMQswCQYDVQQHEwJjZDESMBAGA1UEChMJbG9jYWxob3N0MRIwEAYDVQQLEwls\n" +
                    "b2NhbGhvc3QxEjAQBgNVBAMTCWxvY2FsaG9zdDCBnzANBgkqhkiG9w0BAQEFAAOB\n" +
                    "jQAwgYkCgYEAi5LrPK3cFLME8Agqo4w1ZA/RcYUZ/CE8x8vU/MfWz7DmLn7GEayo\n" +
                    "kbFeYVoyXqdKi2rCPB+ipJw/KNo3WTsrn7eqfSnGxeh1xaPJvQr9lTMRGWhnxoQU\n" +
                    "CtVvNhjzicIdqQoTFZYerxb8RZofzgY3whr6w212Q1VufXcfTbzbJiUCAwEAAaMh\n" +
                    "MB8wHQYDVR0OBBYEFA6H5FhG+9GW4IlIdZvRmklsyXvVMA0GCSqGSIb3DQEBCwUA\n" +
                    "A4GBAH4rvj6Zh++zTG2ikF/zmO++cBUkcyV1bxys2QujaALVCxmLvVCpUaJ1TchS\n" +
                    "IYIJkF3weDCr71z5e7aTmCmsb5aL+li1Ycn4HLvsS5xjI12OkmgDmgjp71vSy70O\n" +
                    "CkLaadSjCnYI6x6lMc7UhuKROXEXKbs1DTxjyAkTnnPkXQeS\n" +
                    "-----END CERTIFICATE-----"));
            iter.insert(14321, cert.get());
            Bytecode result = new Bytecode(constPool);
            result.addInvokestatic("shawnho.ideaagent.Utils", "getRealityResult", "()[B");
            result.addAstore(7);
            iter.insert(24063, result.get());
        }
    }

    private void modifyMybatis(ConstPool constPool, MethodInfo methodInfo, CodeIterator iter) throws BadBytecode {
        if (methodInfo.getName().equals("b") && methodInfo.getDescriptor().equals("()V")) {
            Bytecode bytecode = new Bytecode(constPool);
            bytecode.addOpcode(Opcode.POP);
            bytecode.addLdc(constPool.addStringInfo("https://localhost:7848/q3"));
            iter.insert(16508, bytecode.get());
        } else if (methodInfo.getName().equals("b") && methodInfo.getDescriptor().equals("(Lcom/intellij/util/io/HttpRequests$Request;)V")) {
            Bytecode cert = new Bytecode(constPool);
            cert.addOpcode(Opcode.POP);
            cert.addLdc(constPool.addStringInfo("-----BEGIN CERTIFICATE-----\n" +
                    "MIICYDCCAcmgAwIBAgIEAdWdvDANBgkqhkiG9w0BAQsFADBjMQswCQYDVQQGEwJj\n" +
                    "bjELMAkGA1UECBMCc2MxCzAJBgNVBAcTAmNkMRIwEAYDVQQKEwlsb2NhbGhvc3Qx\n" +
                    "EjAQBgNVBAsTCWxvY2FsaG9zdDESMBAGA1UEAxMJbG9jYWxob3N0MB4XDTE4MDQw\n" +
                    "NjA5MDk0NVoXDTE5MDQwNjA5MDk0NVowYzELMAkGA1UEBhMCY24xCzAJBgNVBAgT\n" +
                    "AnNjMQswCQYDVQQHEwJjZDESMBAGA1UEChMJbG9jYWxob3N0MRIwEAYDVQQLEwls\n" +
                    "b2NhbGhvc3QxEjAQBgNVBAMTCWxvY2FsaG9zdDCBnzANBgkqhkiG9w0BAQEFAAOB\n" +
                    "jQAwgYkCgYEAi5LrPK3cFLME8Agqo4w1ZA/RcYUZ/CE8x8vU/MfWz7DmLn7GEayo\n" +
                    "kbFeYVoyXqdKi2rCPB+ipJw/KNo3WTsrn7eqfSnGxeh1xaPJvQr9lTMRGWhnxoQU\n" +
                    "CtVvNhjzicIdqQoTFZYerxb8RZofzgY3whr6w212Q1VufXcfTbzbJiUCAwEAAaMh\n" +
                    "MB8wHQYDVR0OBBYEFA6H5FhG+9GW4IlIdZvRmklsyXvVMA0GCSqGSIb3DQEBCwUA\n" +
                    "A4GBAH4rvj6Zh++zTG2ikF/zmO++cBUkcyV1bxys2QujaALVCxmLvVCpUaJ1TchS\n" +
                    "IYIJkF3weDCr71z5e7aTmCmsb5aL+li1Ycn4HLvsS5xjI12OkmgDmgjp71vSy70O\n" +
                    "CkLaadSjCnYI6x6lMc7UhuKROXEXKbs1DTxjyAkTnnPkXQeS\n" +
                    "-----END CERTIFICATE-----"));
            iter.insert(14503, cert.get());
            Bytecode result = new Bytecode(constPool);
            result.addInvokestatic("shawnho.ideaagent.Utils", "getRealityResult", "()[B");
            result.addAstore(7);
            iter.insert(24613, result.get());
        }
    }

    @Override
    public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined, ProtectionDomain protectionDomain, byte[] classfileBuffer) throws IllegalClassFormatException {
        if (className == null || (!className.startsWith("com/a/") && !className.startsWith("com/seventh7/"))) {
            return null;
        }
        try {
            if ("com/a/a/a/p".equals(className)) {
                count.incrementAndGet();
            } else if ("com/seventh7/widget/iedis/T".equals(className)) {
                iedis = true;
            } else if ("com/a/b/h/a/b".equals(className)) {
                mybatis = true;
            }
            CtClass ctClass = pool.makeClass(new ByteArrayInputStream(classfileBuffer), false);
            ClassFile classFile = ctClass.getClassFile();
            ConstPool constPool = classFile.getConstPool();
            List methods = classFile.getMethods();
            for (Object method : methods) {
                MethodInfo methodInfo = (MethodInfo) method;
                CodeAttribute codeAttribute = methodInfo.getCodeAttribute();
                if (codeAttribute == null) continue;
                CodeIterator iter = codeAttribute.iterator();
                if (methodInfo.getName().equals("<init>")) {
                    enableAgent(constPool, methodInfo, iter);
                }
                modifyUrl(constPool, methodInfo, iter);
                modifyCert(constPool, methodInfo, iter);
                modifyResult(constPool, methodInfo, iter);
                if ("com/a/a/a/p".equals(className)) {
                    if (count.get() == 1) {
                        if (iedis && !mybatis) {
                            modifyIedis(constPool, methodInfo, iter);
                        } else if (!iedis && mybatis) {
                            modifyMybatis(constPool, methodInfo, iter);
                        } else if (iedis && mybatis) {
                            modifyIedis(constPool, methodInfo, iter);
                        }
                    } else if (count.get() == 2) {
                        modifyMybatis(constPool, methodInfo, iter);
                    }
                }
            }
            return ctClass.toBytecode();
        } catch (Exception e) {
            e.printStackTrace();
        }
        return null;
    }
}
